#ifndef detector_h
#define detector_h

#include "cuda_utils.h"
#include "logging.h"
#include "utils.h"
#include "preprocess.h"
#include "postprocess.h"
#include "model.h"

#include <iostream>
#include <chrono>
#include <cmath>


using namespace nvinfer1;


const static int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1;

class DetectNode
{  
 public:
    std::string engine_name_;
    float cost_time_;
    float kConfThresh;
    float kNmsThresh;
 private:
    Logger gLogger;
        // 选择device
    cudaStream_t stream;
    // Deserialize the engine from file
    IRuntime* runtime = nullptr;
    ICudaEngine* engine = nullptr;
    IExecutionContext* context = nullptr;
    float* gpu_buffers[2];
    float* cpu_output_buffer = nullptr;

 public:
   DetectNode(std::string engine_name, float conf_thresh, float nms_thresh);
   ~DetectNode();
   void prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_output_buffer, float** cpu_output_buffer);
   void infer(IExecutionContext& context, cudaStream_t& stream, void** gpu_buffers, float* output, int batchsize);
   void deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine, IExecutionContext** context);
   void inference(cv::Mat& img);
   void inference(cv::Mat& img, std::vector<Detection>& res);
};









#endif